#include "tree.h"
#include "alignment.h"
#include "MbRandom.h"
#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <istream>
#include <sstream>
#include <fstream>

using namespace std;



Node::Node(void) {

	lft    = NULL;
	rht    = NULL;
	anc    = NULL;
	cls		 = NULL;
	clsPtr = NULL;
	index  = 0;
	clsLength = 0;
	isLeaf = false;
	isChildRoot = false;
	isSubRoot = false;
	flag   = false;
	name   = "";
	
}

Node::~Node(void) {

	if (cls != NULL)
		delete [] cls;
	if (clsPtr != NULL)
		delete [] clsPtr;

}

void Node::initCLS(int nc, int ns) {

	clsLength = nc*ns;
	cls = new double [clsLength];
	clsPtr = new double *[nc];
	if ( !cls )
		{
		cerr << "ERROR: Problem allocating conditional likelihoods" << endl;
		exit(1);
		}
	for (int i=0; i<nc*ns; i++)
		cls[i] = 0.0;
	for (int i=0; i<nc; i++)
		clsPtr[i] = &cls[i * ns];
}

Tree::Tree(Alignment *al, MbRandom *rp, string treeStr) {

	nodes = NULL;
	downPassSequence = NULL;
	alignmentPtr = al;
	ranPtr = rp;
	
	/* set up the conditional likelihoods for each node */
	numChar      = alignmentPtr->getNumChar();
	numStates    = 4;
	
	if (treeStr == "")
		buildRandomTree();
	else
		buildTreeFromNewickString(treeStr);
	
}

Tree::~Tree(void) {

	if (nodes != NULL)
		delete [] nodes;
	if (downPassSequence != NULL)
		delete [] downPassSequence;

}

void Tree::buildRandomTree(void) {

	numTaxa = alignmentPtr->getNumTaxa();
	if (isRooted == true)
		numNodes = 2 * numTaxa;
	else
		numNodes = 2 * numTaxa - 2;

	/* allocate the nodes for the tree */
	nodes = new Node[numNodes];
	for (int i=0; i<numNodes; i++) 
		{
		nodes[i].setIndex(i);
		nodes[i].initCLS(numChar, numStates);
		}
	downPassSequence = new Node*[numNodes];
	for (int i=0; i<numNodes; i++)
			downPassSequence[i] = NULL;
	for (int i=0; i<numTaxa; i++) {
		nodes[i].setName( alignmentPtr->getTaxonName(i) );
		for (int j=0; j<numChar; j++)
			{
			double *clp = nodes[i].getClsPtr(j);
			int nucCode = alignmentPtr->getNucleotide(i, j);
			int nucs[4];
			alignmentPtr->getPossibleNucs(nucCode, nucs);
			for (int k=0; k<numStates; k++)
				{
				if ( nucs[k] == 1 )
					clp[k] = 1.0;
				}
			}
		}

	Node **availableNodes = new Node*[numNodes];
	int numAvailableNodes = 0;
	int nextTip = 0;
	int nextInt = numTaxa;
	
	/* set up three taxon tree */
	int numTaxaInTree = 0;
	Node *p;
	if (isRooted == true)
		p = &nodes[nextInt++];
	else
		{
		p = &nodes[nextTip++];
		numTaxaInTree++;
		}
	root = p;
	Node *q = &nodes[nextInt++];
	availableNodes[numAvailableNodes++] = q;
	p->setLft(q);
	q->setAnc(p);
	p = q;
	q = &nodes[nextTip++];
	numTaxaInTree++;
	availableNodes[numAvailableNodes++] = q;
	p->setLft(q);
	q->setAnc(p);
	q = &nodes[nextTip++];
	numTaxaInTree++;
	availableNodes[numAvailableNodes++] = q;
	p->setRht(q);
	q->setAnc(p);
	
	int numTips = nextTip;
	for (int i=numTips; i<numTaxa; i++)
		{
		p = availableNodes[(int)(ranPtr->uniformRv()*numAvailableNodes)];
		q = p->getAnc();
		Node *a = &nodes[nextTip++];
		numTaxaInTree++;
		Node *b = &nodes[nextInt++];
		if (q->getLft() == p)
			{
			q->setLft(b);
			b->setAnc(q);
			b->setLft(p);
			p->setAnc(b);
			b->setRht(a);
			a->setAnc(b);
			}
		else
			{
			q->setRht(b);
			b->setAnc(q);
			b->setRht(p);
			p->setAnc(b);
			b->setLft(a);
			a->setAnc(b);
			}
		availableNodes[numAvailableNodes++] = a;
		availableNodes[numAvailableNodes++] = b;
		}
	delete [] availableNodes;

	/* initialize the post-order tree traversal sequence */
	getDownPassSequence();
	
	for (int n=0; n<numNodes; n++)
		{
		Node *p = getDownPassNode( n );
		if (p->getAnc() != NULL)
			p->setV( ranPtr->exponentialRv(10.0) );
		else
			p->setV( 0.0 );
		}
#	if 0
	showNodes(root, 3);
	cout << "numTaxaInTree = " << numTaxaInTree << endl;
#	endif
}

/* build a tree up from a string with the Newick description of the tree */
void Tree::buildTreeFromNewickString(string treeStr) {

	/* place all of the tree tokens into a string vector, and count the number of taxa in the tree while we are at it */
	vector<string> words;
	numTaxa = 0;
	int numLeftParantheses = 0;
	string nameStr = "";
	bool readingLength = false;
	bool brlensPresent = false;
	for (int i=0; i<treeStr.length(); i++)
		{
		char ch = treeStr.at(i);
		if ( ch == '(' || ch == ')' || ch == ',' || ch == ':' || ch == ';' )
			{
			string tmp = "";
			tmp += ch;
			words.push_back(tmp);
			nameStr = "";
			if ( ch == ':' )
				{
				readingLength = true;
				brlensPresent = true;
				}
			else if ( ch == '(' )
				numLeftParantheses++;
			//cout << "\"" << ch << "\"" << endl;
			}
		else
			{
			nameStr += ch;
			if ( treeStr.at(i+1) == '(' || treeStr.at(i+1) == ')' || treeStr.at(i+1) == ',' || treeStr.at(i+1) == ':' || treeStr.at(i+1) == ';' )
				{
				words.push_back(nameStr);
				if ( readingLength == false )
					numTaxa++;
				else
					readingLength = false;
				}
			}
		}

	/* decide whether the tree is rooted or unrooted from the number of taxa and the number of left parantheses */
	if ( numLeftParantheses == numTaxa - 1 )
		isRooted = true;
	else if ( numLeftParantheses == numTaxa - 2 )
		isRooted = false;
	else
		{
		cerr << "Inconsistent number of parantheses in tree description" << endl;
		exit(1);
		}
		
	/* calculate the number of nodes in the tree */
	if (isRooted == true)
		numNodes = 2 * numTaxa;
	else
		numNodes = 2 * numTaxa - 2;
		
	/* check for consistency in the number of taxa between the data and the tree */
	if ( numTaxa != alignmentPtr->getNumTaxa() )
		{
		cerr << "Problem reading the correct number of taxa." << endl;
		exit(1);
		}

	/* allocate the nodes for the tree */
	nodes = new Node[numNodes];
	for (int i=0; i<numNodes; i++)
		{
		nodes[i].setIndex(i);
		nodes[i].initCLS(numChar, numStates);
		}
	downPassSequence = new Node*[numNodes];
	for (int i=0; i<numNodes; i++)
			downPassSequence[i] = NULL;
	/* build the tree */
	int nextIntNode = numTaxa;
	readingLength = false;
	Node *p, *q;
	for (vector<string>::iterator s=words.begin(); s != words.end(); s++)
		{
		if ( (*s) == "(" )
			{
			/* add an internal node */
			if (nextIntNode == numTaxa)
				{
				p = &nodes[nextIntNode++];
				root = p;
				}
			else
				{
				if (p->getLft() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setAnc(p);
					p->setLft(q);
					p = q;
					}
				else if (p->getRht() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setAnc(p);
					p->setRht(q);
					p = q;
					}
				else if (p->getAnc() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setLft(p);
					p->setAnc(q);
					q->setFlag( true );
					p = q;
					root = p;
					}
				else
					{
					cout << "Problem finding an open node" << endl;
					exit(1);
					}
				}
			}
		else if ( (*s) == ")" )
			{
			/* go down */
			if (p->getFlag() == false)
				p = p->getAnc();
			else
				p = p->getLft();
			}
		else if ( (*s) == "," )
			{
			/* go down */
			if (p->getFlag() == false)
				p = p->getAnc();
			else
				p = p->getLft();
			}
		else if ( (*s) == ":" )
			{
			/* start reading a digit */
			readingLength = true;
			}
		else if ( (*s) == ";" )
			{
			/* end of the tree */

			}
		else
			{
			if ( readingLength == true )
				{
				/* add a branch length */
				double x;
				istringstream buf( (*s) );
				buf >> x;
				//p->setV( x );
				readingLength = false;
				}
			else
				{
				/* add a tip node */
				int taxonIndex = alignmentPtr->getTaxonIndex( (*s) );
				if (taxonIndex == -1)
					{
					cerr << "Couldn't find \"" << (*s) << "\" in the list of taxon names." << endl;
					exit(1);
					}
				
				if (p->getLft() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setAnc(p);
					p->setLft(q);
					p = q;
					}
				else if (p->getRht() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setAnc(p);
					p->setRht(q);
					p = q;
					}
				else if (p->getAnc() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setLft(p);
					p->setAnc(q);
					q->setFlag( true );
					p = q;
					root = p;
					}
				else
					{
					cerr << "Problem finding an open node!" << endl;
					exit(1);
					}
				p->setIsLeaf( true );
				p->setName( *s );
				for (int j=0; j<numChar; j++)
					{
					double *clp = p->getClsPtr(j);
					int nucCode = alignmentPtr->getNucleotide(p->getIndex(), j);
					int nucs[4];
					alignmentPtr->getPossibleNucs(nucCode, nucs);
					for (int k=0; k<numStates; k++)
						{
						if ( nucs[k] == 1 )
							clp[k] = 1.0;
						}
					}
				}
			}
		}
		
	if (isRooted == true)
		{
		p = &nodes[nextIntNode];
		q = root;
		p->setLft( q );
		q->setAnc( p );
		root = p;
		}
		
	/* initialize the post-order tree traversal sequence */
	getDownPassSequence();
	
	for (int n=0; n<numNodes; n++)
		{
		Node *p = getDownPassNode( n );
		if (p->getAnc() != NULL)
			p->setV( ranPtr->exponentialRv(10.0) );
		else
			p->setV( 0.0 );
		}

}

/* get the post-order tree traversal sequence */
void Tree::getDownPassSequence(void) {

	int i = 0;
	passDn(root, &i);

}

/* recursive function filling in the post-order traversal */
void Tree::passDn(Node *p, int *x) {

	if (p != NULL)
		{
		passDn(p->getLft(), x);
		passDn(p->getRht(), x);
		downPassSequence[(*x)++] = p;
		}
		
}

/* print out the tree in an indented format */
void Tree::print(void) {

	showNodes(root, 3);

}

/* show the nodes of the tree in an indented format */
void Tree::showNodes(Node *p, int indent) {

	if (p != NULL)
		{
		for (int i=0; i<indent; i++)
			cout << " ";
		cout << dex(p) << " (" << dex(p->getLft()) << ", " << dex(p->getRht()) << ", " << dex(p->getAnc()) << ") ";
		if ( p != root )
			cout << fixed << setprecision(6) << p->getV();
		if (p->getIsLeaf() == true)
			cout << " (" << p->getName() << ")";
		if (p == root)
			cout << " <- Root" << endl;
		else
			cout << endl;
		showNodes (p->getLft(),  indent + 2);
		showNodes (p->getRht(), indent + 2);
		}
   
}

/* a small function that returns the index of the node (or -1 if the node is NULL) */
int Tree::dex(Node *p) {

	return ( p == NULL ? -1 : p->getIndex() );

}


SubTree::SubTree(Tree *tr, Node *rt, int idx) {
	
	subRoot = rt;
	subRoot->setIsSubRoot(true);
	index = idx;
	numNodesST = 0;
	numChildRoot = 0;
	downPassSequenceST = NULL;
	ancSubTree = NULL;
	clsST = NULL;
	clsPtrST = NULL;
	setNumNodesST(subRoot, &numNodesST);
	cout << numNodesST << endl;
	
	downPassSequenceST = new Node*[numNodesST];
	for (int i=0; i<numNodesST; i++)
		downPassSequenceST[i] = NULL;
	getDownPassSequenceST();
	
	if (subRoot->getAnc() != NULL) 
		subRoot->setIsChildRoot(true);
	
}

SubTree::~SubTree(void) {
	
	if (downPassSequenceST != NULL)
		delete [] downPassSequenceST;
	if (clsST != NULL)
		delete [] clsST;
	if (clsPtrST != NULL)
		delete [] clsPtrST;

}

/* get the post-order tree traversal sequence */
void SubTree::getDownPassSequenceST(void) {

	int i = 0;
	passDnST(subRoot, &i);

}

/* recursive function filling in the post-order traversal */
void SubTree::passDnST(Node *p, int *x) {

	if (p != NULL)
		{
			if (p->getIsChildRoot())
				{
					numChildRoot++;
					downPassSequenceST[(*x)++] = p;
					return;
				}
		passDnST(p->getLft(), x);
		passDnST(p->getRht(), x);
		downPassSequenceST[(*x)++] = p;		
		}
		
}

/* recursively calculate the total number of nodes in tree */
void SubTree::setNumNodesST(Node *sr, int *nNST) {

	if (sr != NULL)
		{
			if (sr->getIsChildRoot())
				{
					(*nNST)++;
					return;
				}
		setNumNodesST(sr->getLft(), nNST);
		setNumNodesST(sr->getRht(), nNST);
		(*nNST)++;		
		}
		
}

/* gather cls from all nodes in subtree */
double* SubTree::getClsST(void) {
	int nodeClsLength = subRoot->getClsLength();
	double *nodeCls = new double [nodeClsLength];
	double *clsST = new double[numNodesST * nodeClsLength];
	
	for (int n=0; n<numNodesST; n++)
		{
		nodeCls = downPassSequenceST[n]->getCls();
		for (int j=0; j<nodeClsLength; j++)
			clsST[n * nodeClsLength + j] = nodeCls[j];				
		}
	
	return clsST;
		
}

// Given a cls array returns a clsPtr for the subtree
double** SubTree::getClsPtrST(double *CST) {
	
	int nodeClsLength = subRoot->getClsLength();
	double *nodeCls = new double [nodeClsLength];
	double **clsPtrST = new double*[numNodesST];
	
	for (int n=0; n<numNodesST; n++)
			clsPtrST[n] = &CST[n * nodeClsLength];
	
	return clsPtrST;
		
}

// Given a cls array sets the node cls
void SubTree::setClsST(double *CST) {
	int nodeClsLength = subRoot->getClsLength();
	double *nodeCls = new double [nodeClsLength];
	
	for (int n=0; n<numNodesST; n++)
		{
		Node *p = downPassSequenceST[n];
		for (int j=0; j<nodeClsLength; j++)
			p->setCls(j, CST[n * nodeClsLength + j]);				
		}
		
}

Queue::Queue(int numST) {
	
	nextQueue = new SubTree*[numST];
	queueStart = 0;
	queueEnd = 0;

	for (int i=0; i<numST; i++)
			nextQueue[i] = NULL;

}

Queue::~Queue(void) {
	
	delete [] nextQueue;

}

void Queue::completeSubTree(SubTree *st) {
	SubTree *ancST = st->getAncSubTree();
	if (ancST != NULL)
		{
			ancST->numChildRoot--;
			if(ancST->numChildRoot == 0)
				enqueue(ancST);
		}
}

void Queue::enqueue(SubTree *st) {
	nextQueue[queueEnd] = st;
	queueEnd++;
}

SubTree* Queue::dequeue(void) {
	if(queueEnd != queueStart)
		{
			return nextQueue[queueStart++];
		}
	return NULL;
}

